﻿using System;
using System.Collections.Generic;
using System.Text;
using Osteox.ManagedCv;
using System.Runtime.InteropServices;

namespace watershed
{
    class Program
    {
        static ManagedImage marker_mask = null;
        static ManagedImage markers = null;
        static ManagedImage img0 = null;
        static ManagedImage img = null;
        static ManagedImage img_gray = null;
        static ManagedImage wshed = null;
        static Point2D32s prev_pt = new Point2D32s(-1, -1);

        static void on_mouse(WindowEventId eventid, int x, int y, WindowEventFlag flags, IntPtr param)
        {
            if (ManagedImage.IsNullOrInvalid(img))
                return;

            if (eventid == WindowEventId.LButtonUp || (flags & WindowEventFlag.FlagLButton) == WindowEventFlag.None)
                prev_pt = new Point2D32s(-1, -1);
            else if (eventid == WindowEventId.LButtonDown)
                prev_pt = new Point2D32s(x, y);
            else if (eventid == WindowEventId.MouseMove && (flags & WindowEventFlag.FlagLButton) != WindowEventFlag.None)
            {
                Point2D32s pt = new Point2D32s(x, y);
                if (prev_pt.X < 0)
                    prev_pt = pt;
                ManagedCvInvoke.Line(marker_mask, prev_pt, pt, ManagedCvInvoke.ScalarAll(255), 5, LineTypeFlag.Connected8, 0);
                ManagedCvInvoke.Line(img, prev_pt, pt, ManagedCvInvoke.ScalarAll(255), 5, LineTypeFlag.Connected8, 0);
                prev_pt = pt;
                ManagedCvInvoke.ShowImage("image", img);
            }
        }

        static void Main(string[] args)
        {
            string filename = args.Length >= 2 ? args[1] : "fruits.jpg";
            Rng rng = ManagedCvInvoke.RNG(-1);
            img0 = ManagedCvInvoke.LoadImageW(filename);

            if (ManagedImage.IsNullOrInvalid(img0))
                return;

            Console.Write("Hot keys: \n");
            Console.Write("\tESC - quit the program\n");
            Console.Write("\tr - restore the original image\n");
            Console.Write("\tw or SPACE - run watershed algorithm\n");
            Console.Write("\t\t(before running it, roughly mark the areas on the image)\n");
            Console.Write("\t  (before that, roughly outline several markers on the image)\n");

            ManagedCvInvoke.NamedWindow("image", WindowMode.Autosize);
            ManagedCvInvoke.NamedWindow("watershed transform", WindowMode.Autosize);

            img = ManagedCvInvoke.CloneImage(img0);
            img_gray = ManagedCvInvoke.CloneImage(img0);
            wshed = ManagedCvInvoke.CloneImage(img0);
            marker_mask = ManagedCvInvoke.CreateImage(ManagedCvInvoke.GetSize(img), PixelDepth.Depth8U, 1);
            markers = ManagedCvInvoke.CreateImage(ManagedCvInvoke.GetSize(img), PixelDepth.Depth32S, 1);
            ManagedCvInvoke.CvtColor(img, marker_mask, ColorMapType.BGR2GRAY);
            ManagedCvInvoke.CvtColor(marker_mask, img_gray, ColorMapType.GRAY2BGR);

            ManagedCvInvoke.Zero(marker_mask);
            ManagedCvInvoke.Zero(wshed);
            ManagedCvInvoke.ShowImage("image", img);
            ManagedCvInvoke.ShowImage("watershed transform", wshed);
            ManagedCvInvoke.SetMouseCallback("image", on_mouse, IntPtr.Zero);

            for (; ; )
            {
                int c = ManagedCvInvoke.WaitKey(0);

                if ((char)c == 27)
                    break;

                if ((char)c == 'r')
                {
                    ManagedCvInvoke.Zero(marker_mask);
                    ManagedCvInvoke.Copy(img0, img);
                    ManagedCvInvoke.ShowImage("image", img);
                }

                if ((char)c == 'w' || (char)c == ' ')
                {
                    ManagedMemStorage storage = ManagedCvInvoke.CreateMemStorage(0);
                    ManagedSeq contours = null;
                    ManagedMat color_tab = null;
                    int i, j, comp_count = 0;
                    ManagedCvInvoke.FindContours(marker_mask, storage, out contours,
                        Marshal.SizeOf(typeof(Contour)), ContourRetrMode.CComp, ContourApproxMethod.Simple);
                    ManagedCvInvoke.Zero(markers);

                    for (; contours != null && contours.IsValid; contours = contours.HNext, comp_count++)
                    {
                        ManagedCvInvoke.DrawContours(markers, contours,
                            ManagedCvInvoke.ScalarAll(comp_count + 1),
                            ManagedCvInvoke.ScalarAll(comp_count + 1),
                            -1,
                            -1,
                            LineTypeFlag.Connected8,
                            ManagedCvInvoke.Point(0, 0));
                    }

                    color_tab = ManagedCvInvoke.CreateMat(1, comp_count, MatFlag.CV8UC3);
                    ManagedArrPixelAccessorBgr8U colormapAccessor = new ManagedArrPixelAccessorBgr8U(color_tab);
                    for (i = 0; i < comp_count; i++)
                    {
                        colormapAccessor[0, i] = new Bgr<byte>(
                            (byte)(ManagedCvInvoke.RandInt(ref rng) % 180 + 50),
                            (byte)(ManagedCvInvoke.RandInt(ref rng) % 180 + 50),
                            (byte)(ManagedCvInvoke.RandInt(ref rng) % 180 + 50));
                    }

                    double t = (double)ManagedCvInvoke.GetTickCount();
                    ManagedCvInvoke.Watershed(img0, markers);
                    t = (double)ManagedCvInvoke.GetTickCount() - t;
                    Console.Write("exec time = {0}ms\n", t / (ManagedCvInvoke.GetTickFrequency() * 1000.0));

                    ManagedArrPixelAccessorGray32S markersAccessor = new ManagedArrPixelAccessorGray32S(markers);
                    ManagedArrPixelAccessorBgr8U wshedAccessor = new ManagedArrPixelAccessorBgr8U(wshed);
                    Bgr<byte> pixel = new Bgr<byte>();
                    // paint the watershed image
                    for (i = 0; i < markers.Height; i++)
                    {
                        for (j = 0; j < markers.Width; j++)
                        {
                            int idx = markersAccessor[i, j].Intensity;
                            if (idx == -1)
                            {
                                pixel.B = pixel.G = pixel.R = 255;
                            }
                            else if (idx <= 0 || idx > comp_count)
                            {
                                pixel.B = pixel.G = pixel.R = 0;
                            }
                            else
                            {
                                pixel = colormapAccessor[0, idx - 1];
                            }
                            wshedAccessor[i, j] = pixel;
                        }
                    }
                    ManagedCvInvoke.AddWeighted(wshed, 0.5, img_gray, 0.5, 0, wshed);
                    ManagedCvInvoke.ShowImage("watershed transform", wshed);
                }
            }
        }
    }
}
